// onnx_to_cpp.C

void onnx_to_cpp()
{
   using namespace TMVA::Experimental;
   SOFIE::RModelParser_ONNX parser;
   SOFIE::RModel model = parser.Parse("./VRlL_real_500k_evts_model.onnx");
   model.SetOptimizationLevel(SOFIE::OptimizationLevel::kBasic);
   model.Generate();
   model.PrintRequiredInputTensors();

   model.OutputGenerated("./VRlL_real_500k_evts_model.hxx");
}

// sofie_ad.C

#include "VRlL_real_500k_evts_model.hxx"

#include <Math/CladDerivator.h>

using Sess = TMVA_SOFIE_VRlL_real_500k_evts_model::Session;

// Wrapper functions for Clad
float my_func(Sess const &session, float const *tensor_x, float *tensor_theory_params)
{
   float out = 0.;
   TMVA_SOFIE_VRlL_real_500k_evts_model::doInfer(session, tensor_x, tensor_theory_params, &out);
   return out;
}

float my_func_wrapper(Sess const &session, float const *tensor_x, float *tensor_theory_params)
{
   return my_func(session, tensor_x, tensor_theory_params);
}

void sofie_ad()
{
   // Let's go Clad!
   clad::gradient(my_func_wrapper, "tensor_theory_params");

   // Get a function pointer to the pullback. If you are unsure what the
   // signature is, try to cast the pullback to some function pointer, like
   // static_cast<void (*)(float)>(my_func_pullback) in the interpreter, and
   // the compiler will tell you what the real signature is.
   using Grad_t = void (*)(const Sess &, const float *, float *, float, Sess *, float *);

   // Get the functions from the interpreter (remove semicolor to get the code printed)
   auto grad = reinterpret_cast<Grad_t>(gInterpreter->ProcessLine("my_func_pullback;"));

   std::vector<float> input1{5.0, 2.0, 1.0, -1.0, 1.0};
   std::vector<float> input2{0.0};

   // A trick: pre-allocate session struct both for the forward pass and
   // backward pass, to that no memory allocation of intermediate tensors has
   // to happend in the gradient.
   TMVA_SOFIE_VRlL_real_500k_evts_model::Session s("VRlL_real_500k_evts_model.dat");
   TMVA_SOFIE_VRlL_real_500k_evts_model::Session d_s("VRlL_real_500k_evts_model.dat");

   // Calculate numerical gradient
   auto numDiff = [&](int i) {
      const float eps = 1e-4;
      std::vector<float> p{input2};
      p[i] = input2[i] - eps;
      float funcValDown = my_func(s, input1.data(), p.data());
      p[i] = input2[i] + eps;
      float funcValUp = my_func(s, input1.data(), p.data());
      return (funcValUp - funcValDown) / (2 * eps);
   };

   for (std::size_t i = 0; i < input2.size(); ++i) {
      std::cout << i << ":" << std::endl;
      std::cout << "  numr : " << numDiff(i) << std::endl;
   }

   // Calculate gradient with Clad
   float grad_output[]{0., 0., 0., 0., 0.};
   grad(s, input1.data(), input2.data(), 1.0, &d_s, grad_output);

   std::cout << "  clad : " << grad_output[0] << std::endl;
}